import numpy as np
import matplotlib.pyplot as plt
import pylab

coordinates = np.array([
    [6, 25], [7, 43], [9, 56], [10, 70], [11, 28],
    [12, 17], [12, 38], [15, 5], [15, 14], [15, 56],
    [16, 19], [17, 64], [20, 30], [21, 48], [21, 45],
    [21, 36], [22, 53], [22, 22], [26, 29], [26, 13],
    [26, 59], [27, 24], [29, 39], [30, 50], [30, 20],
    [30, 60], [31, 76], [33, 34], [33, 44], [35, 51],
    [35, 16], [35, 60], [36, 6], [36, 26], [38, 33],
    [40, 37], [40, 66], [40, 60], [40, 20], [41, 46],
    [43, 26], [44, 13], [45, 42], [45, 35], [47, 66],
    [48, 21], [50, 30], [50, 40], [50, 50], [50, 70],
    [50, 4], [50, 15], [51, 42], [52, 26], [54, 38],
    [54, 10], [55, 34], [55, 45], [55, 50], [55, 65],
    [55, 57], [55, 20], [57, 72], [59, 5], [60, 15],
    [62, 57], [62, 48], [62, 35], [62, 24], [64, 4],
    [65, 27], [66, 14], [66, 8], [67, 41], [70, 64]
])


def getdistmat(coordinates):
    num = coordinates.shape[0]
    distmat = np.zeros((coordinates.shape[0], coordinates.shape[0]))
    for i in range(num):
        for j in range(i, num):
            distmat[i][j] = distmat[j][i] = np.linalg.norm(coordinates[i] - coordinates[j])
    return distmat


distmat = getdistmat(coordinates)
numant = 40  # 蚂蚁个数
numcity = coordinates.shape[0]  # 城市个数
alpha = 1  # 信息素重要程度因子
beta = 5  # 启发函数重要程度因子
rho = 0.1  # 信息素的挥发速度
Q = 1
iter = 0
itermax = 250
etatable = 1.0 / (distmat + np.diag([1e10] * numcity))  # 启发函数矩阵，表示蚂蚁从城市i转移到矩阵j的期望程度
pheromonetable = np.ones((numcity, numcity))  # 信息素矩阵
pathtable = np.zeros((numant, numcity)).astype(int)  # 路径记录表
distmat = getdistmat(coordinates)  # 城市的距离矩阵
lengthaver = np.zeros(itermax)  # 各代路径的平均长度
lengthbest = np.zeros(itermax)  # 各代及其之前遇到的最佳路径长度
pathbest = np.zeros((itermax, numcity))  # 各代及其之前遇到的最佳路径长度

while iter < itermax:
    if numant <= numcity:
        pathtable[:, 0] = np.random.permutation(range(0, numcity))[:numant]
    else:
        pathtable[:numcity, 0] = np.random.permutation(range(0, numcity))[:]
        pathtable[numcity:, 0] = np.random.permutation(range(0, numcity))[:numant - numcity]
    length = np.zeros(numant)
    for i in range(numant):
        visiting = pathtable[i, 0]
        unvisited = set(range(numcity))
        unvisited.remove(visiting)
        for j in range(1, numcity):
            listunvisited = list(unvisited)
            probtrans = np.zeros(len(listunvisited))
            for k in range(len(listunvisited)):
                probtrans[k] = np.power(pheromonetable[visiting][listunvisited[k]], alpha) \
                               * np.power(etatable[visiting][listunvisited[k]], alpha)
            cumsumprobtrans = (probtrans / sum(probtrans)).cumsum()
            cumsumprobtrans -= np.random.rand()
            k = listunvisited[(np.where(cumsumprobtrans > 0)[0])[0]]

            pathtable[i, j] = k
            unvisited.remove(k)
            length[i] += distmat[visiting][k]
            visiting = k
        length[i] += distmat[visiting][pathtable[i, 0]]
    lengthaver[iter] = length.mean()
    if iter == 0:
        lengthbest[iter] = length.min()
        pathbest[iter] = pathtable[length.argmin()].copy()
    else:
        if length.min() > lengthbest[iter - 1]:
            lengthbest[iter] = lengthbest[iter - 1]
            pathbest[iter] = pathbest[iter - 1].copy()
        else:
            lengthbest[iter] = length.min()
            pathbest[iter] = pathtable[length.argmin()].copy()
    # 更新信息素
    changepheromonetable = np.zeros((numcity, numcity))
    for i in range(numant):
        for j in range(numcity - 1):
            changepheromonetable[pathtable[i, j]][pathtable[i, j + 1]] += Q / distmat[pathtable[i, j]][
                pathtable[i, j + 1]]
        changepheromonetable[pathtable[i, j + 1]][pathtable[i, 0]] += Q / distmat[pathtable[i, j + 1]][pathtable[i, 0]]
    pheromonetable = (1 - rho) * pheromonetable + changepheromonetable
    iter += 1
    print("iter:", iter)


fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(12, 10))
axes[0].plot(lengthaver, 'k', marker=u'')
axes[0].set_title('Average Length')
axes[0].set_xlabel(u'iteration')

axes[1].plot(lengthbest, 'k', marker=u'')
axes[1].set_title('Best Length')
axes[1].set_xlabel(u'iteration')
fig.savefig('average_best.png', dpi=500, bbox_inches='tight')
plt.show()
print(lengthbest[-1])


bestpath = pathbest[-1]
plt.plot(coordinates[:, 0], coordinates[:, 1], 'r.', marker=u'$\cdot$')
# plt.xlim([0, 1])
# plt.ylim([0, 1])

for i in range(numcity - 1):
    m = int(bestpath[i])
    n = int(bestpath[i + 1])
    plt.plot([coordinates[m][0], coordinates[n][0]], [coordinates[m][1], coordinates[n][1]], 'k')
plt.plot([coordinates[int(bestpath[0])][0], coordinates[int(n)][0]],
         [coordinates[int(bestpath[0])][1], coordinates[int(n)][1]], 'b')
ax = plt.gca()
ax.set_title("Best Path, length is % s" % lengthbest[-1])
ax.set_xlabel('X axis')
ax.set_ylabel('Y_axis')

plt.savefig('best path.png', dpi=500, bbox_inches='tight')
plt.show()
